import os
import rclpy
from rclpy.node import Node
from cv_bridge import CvBridge
import cv2
from sensor_msgs.msg import Image
from semantic_nav.bag_utils import BagUtilsBase, BagTimeUtils
from semantic_nav.log_utils import get_logger
# Replace these with your actual frames, and topics
IMAGE_TOPIC = "/camera/color/image_raw"
logger = get_logger()


def play_video(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        logger.error("Error opening video file")
    # Get the fps and resolution of video
    fps = cap.get(cv2.CAP_PROP_FPS)
    if fps == 0:
        logger.error("Error getting FPS of video")
        return
    frame_interval = int(1000 / fps)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    # Create a window to show video
    cv2.namedWindow("Video", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Video", width, height)
    # Read and show each frame image cyclically
    while True:
        ret, frame = cap.read()
        if not ret:
            logger.debug("Video ended")
            break
        cv2.imshow("Video", frame)
        # Press key 'q' to exit
        if cv2.waitKey(frame_interval) & 0xFF == ord('q'):
            break
    # Release resouces
    cap.release()
    cv2.destroyAllWindows()


class Bag2Video(BagUtilsBase, Node):
    def __init__(self, bag_path, output_path, image_topic=IMAGE_TOPIC):
        super().__init__(bag_path)
        super(BagUtilsBase, self).__init__("bag_to_video")
        self.declare_parameter("image_topic", image_topic)
        self.topic_name = None
        self.cv_bridge = CvBridge()
        self.video_path = os.path.join(output_path, 'output_rgb.mp4')
        self.fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.bag_time_utils = BagTimeUtils(bag_path)
        self.video_writer = None

    def callback(self, topic, image_msg, timestamp):
        if topic == self.topic_name:
            if self.video_writer is None:
                self.video_writer = cv2.VideoWriter(self.video_path, self.fourcc, self.frequency, (image_msg.width, image_msg.height))
            cv_image = self.cv_bridge.imgmsg_to_cv2(image_msg, desired_encoding="bgr8")
            self.video_writer.write(cv_image)
        return True

    def write_video(self):
        self.topic_name = self.get_parameter("image_topic").value
        if not self.has_topic(self.topic_name):
            logger.debug(f"Topic '{self.topic_name}' not found in bag file.")
            return
        if self.get_msg_type(self.topic_name) != Image:
            logger.debug(f"Topic '{self.topic_name}' is not of type Image.")
            return
        self.frequency = round(self.bag_time_utils.get_frequency(self.topic_name))
        logger.debug(f"Frequency of topic '{self.topic_name}' is {self.frequency} Hz")
        self.process()
        self.video_writer.release()
        logger.debug("Finished processing bag file and storing image in video")
        play_video(self.video_path)


def main(args=None):
    rclpy.init(args=args)
    from ament_index_python.packages import get_package_prefix
    package_path = get_package_prefix('semantic_nav')
    bag_path = os.path.join(package_path, "bag", "test_bag")
    output_path = os.path.join(package_path, "video")
    bag2video = Bag2Video(bag_path, output_path)
    bag2video.write_video()

    rclpy.shutdown()


if __name__ == "__main__":
    main()
